import time
import cv2
import imagehash
from PIL import Image
import os
import traceback
import numpy as np
import cv2
import hashlib
from concurrent.futures import ThreadPoolExecutor
from moviepy.video.io.VideoFileClip import VideoFileClip
from dataclasses import dataclass
import cupy as cp

VIDEO_HASH_THRESHOLD = 0.0005

@dataclass
class VideoFile:
    path: str
    width: int
    height: int
    fps: float
    duration: float
    hash: int


def get_video_info(video_path):
    # start_time = time.perf_counter() * 1000
    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # print(f"1耗时: {time.perf_counter() * 1000 - start_time:.2f} 毫秒")

    # start_time = time.perf_counter() * 1000
    # duration = get_video_duration(video_path)
    duration = frames / fps if fps > 0 else 0
    # print(f"2耗时: {time.perf_counter() * 1000 - start_time:.2f} 毫秒")

    # start_time = time.perf_counter() * 1000
    hash = compute_video_hash(cap)
    # print(f"3耗时: {time.perf_counter() * 1000 - start_time:.2f} 毫秒")

    cap.release()
    return VideoFile(video_path, width, height, fps, duration, hash)


def get_video_duration(file_path):
    with VideoFileClip(file_path) as video:
        return video.duration
    
def is_video_hash_close(a, b):
    if a == 0 and b == 0:
        return True
    max_val = max(abs(a), abs(b))
    if max_val == 0:
        return False  # 避免除以零
    relative_diff = abs(a - b) / max_val
    return relative_diff < VIDEO_HASH_THRESHOLD

def compute_video_hash(cap, sample_frames=2):
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    hash_values = []
    
    for i in range(sample_frames):
        pos = int(frame_count/(sample_frames+1)*(i+1))
        cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
        ret, frame = cap.read()
        if ret:
            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            hash_obj = imagehash.average_hash(pil_img)
            hash_values.append(int(str(hash_obj), 16))  # 将哈希转为16进制整数
    return sum(hash_values)//len(hash_values) if hash_values else None

def compute_video_hash_gpu(cap, sample_frames=2):
    # 使用CUDA加速的视频解码器
    cap = cv2.cudacodec.createVideoReader(cap.get(cv2.CAP_PROP_POS_AVI_RATIO))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    hash_values = []
    
    for i in range(sample_frames):
        pos = int(frame_count/(sample_frames+1)*(i+1))
        cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
        ret, gpu_frame = cap.read()
        
        if ret:
            # 将GPU帧转换为CPU内存中的numpy数组
            frame = gpu_frame.download()
            # 使用cupy加速颜色空间转换
            gpu_array = cp.asarray(frame)
            gpu_rgb = cp.zeros_like(gpu_array)
            cp.cuda.runtime.cvtColor(gpu_array, gpu_rgb, cv2.COLOR_BGR2RGB)
            
            # 转换回PIL图像
            rgb_array = cp.asnumpy(gpu_rgb)
            pil_img = Image.fromarray(rgb_array)
            
            hash_obj = imagehash.average_hash(pil_img)
            hash_values.append(int(str(hash_obj), 16))
            
    return sum(hash_values)//len(hash_values) if hash_values else None